import gradio as gr
import pickle
import numpy as np
from PIL import Image

# 加载保存的KNN模型
def load_model(model_path="best_knn_model.pkl"):
    """加载保存的KNN模型"""
    try:
        with open(model_path, 'rb') as f:
            model = pickle.load(f)
        print("模型加载成功")
        return model
    except Exception as e:
        print(f"模型加载失败: {e}")
        return None

# 预处理图像以适应模型输入
def preprocess_image(img):
    """
    预处理图像：将图像转为28x28的灰度图，
    并将其转换为模型所需的输入格式
    """
    # 将图像转为PIL Image对象（如果不是的话）
    if not isinstance(img, Image.Image):
        img = Image.fromarray(img)
    
    # 转为灰度图并调整大小为28x28
    img = img.convert('L')  # 转为灰度图
    img = img.resize((28, 28), Image.LANCZOS)  # 调整大小
    
    # 转换为numpy数组并归一化
    img_array = np.array(img)
    img_array = img_array / 255.0  # 归一化到0-1范围
    
    # 展平为一维数组（784个特征）
    img_flattened = img_array.flatten()
    
    # 调整形状以适应模型输入
    return img_flattened.reshape(1, -1)

# 预测函数
def predict_digit(img):
    """预测手写数字"""
    if model is None:
        return "模型加载失败，请检查模型文件"
    
    try:
        # 预处理图像
        processed_img = preprocess_image(img)
        
        # 进行预测
        prediction = model.predict(processed_img)
        
        # 返回预测结果
        return f"预测结果: {prediction[0]}"
    except Exception as e:
        return f"预测出错: {str(e)}"

# 加载模型
model = load_model()

# 创建Gradio界面
with gr.Blocks(title="手写数字识别") as demo:
    gr.Markdown("# 手写数字识别")
    gr.Markdown("在下方的画布上写下一个数字（0-9），然后点击预测按钮查看结果")
    
    with gr.Row():
        # 创建绘图组件，指定画布大小为200x200，背景为白色
        sketchpad = gr.Sketchpad(height=200, width=200, label="请在此处手写数字")
        
        # 输出结果组件
        output = gr.Textbox(label="预测结果")
    
    # 预测按钮
    predict_btn = gr.Button("预测")
    
    # 设置按钮点击事件
    predict_btn.click(
        fn=predict_digit,
        inputs=sketchpad,
        outputs=output
    )
    
    # 清除按钮
    clear_btn = gr.Button("清除")
    clear_btn.click(
        fn=lambda: None,
        inputs=None,
        outputs=sketchpad,
        queue=False
    )

# 启动应用
if __name__ == "__main__":
    demo.launch()